"""
基金相关的Celery任务
"""
import logging
import datetime
from celery import group
from celery.utils.log import get_task_logger
from sqlalchemy import func

from modules.tasks.celery_app import app
from modules.models.base import get_db
from modules.models.stock import SecurityInfo, CollectionTask
from modules.data_collector import collector

# 配置任务日志
logger = get_task_logger(__name__)


@app.task(bind=True, queue='etf_data', max_retries=3, retry_backoff=60)
def update_fund_list(self, fund_type="all"):
    """
    更新基金列表信息
    
    Args:
        fund_type: 基金类型，可选值：etf/lof/stock/bond/mix/money/qdii/index/all
    """
    logger.info(f"开始更新{fund_type}基金列表信息")
    try:
        new_count, update_count = collector.collect_market_data(
            "fund_list", 
            fund_type=fund_type
        )
        logger.info(f"{fund_type}基金列表更新完成: 新增 {new_count} 只，更新 {update_count} 只")
        return {
            "status": "success",
            "fund_type": fund_type,
            "new_count": new_count,
            "update_count": update_count
        }
    except Exception as e:
        logger.error(f"更新{fund_type}基金列表失败: {str(e)}")
        raise self.retry(exc=e)


@app.task(bind=True, queue='etf_data', max_retries=3, retry_backoff=60)
def collect_fund_nav_data(self, fund_code, start_date=None, end_date=None):
    """
    采集单只基金的净值数据
    
    Args:
        fund_code: 基金代码
        start_date: 开始日期, 格式为YYYYMMDD
        end_date: 结束日期, 格式为YYYYMMDD
    """
    logger.info(f"开始采集基金 {fund_code} 的净值数据")
    
    try:
        # 创建任务记录
        db = next(get_db())
        task_record = CollectionTask(
            task_id=self.request.id,
            task_type='fund_nav',
            security_code=fund_code,
            start_date=datetime.datetime.strptime(start_date, '%Y%m%d').date() if start_date else None,
            end_date=datetime.datetime.strptime(end_date, '%Y%m%d').date() if end_date else None,
            status='pending'
        )
        db.add(task_record)
        db.commit()
        db.close()
        
        # 执行数据采集
        count = collector.collect_market_data(
            "fund_nav", 
            code=fund_code, 
            start_date=start_date, 
            end_date=end_date, 
            task_id=self.request.id
        )
        
        logger.info(f"基金 {fund_code} 净值数据采集完成，获取 {count} 条数据")
        
        # 如果有新数据，触发技术指标计算
        if count > 0:
            # 使用链式任务计算技术指标
            from modules.tasks.stock_tasks import calculate_stock_indicators
            calculate_stock_indicators.delay(fund_code)
        
        return {
            "status": "success",
            "fund_code": fund_code,
            "count": count
        }
    except Exception as e:
        logger.error(f"采集基金 {fund_code} 净值数据失败: {str(e)}")
        # 更新任务记录状态
        try:
            db = next(get_db())
            task = db.query(CollectionTask).filter(CollectionTask.task_id == self.request.id).first()
            if task:
                task.status = 'failed'
                task.error_message = str(e)
                task.update_time = datetime.datetime.now()
                db.commit()
        except Exception as db_e:
            logger.error(f"更新任务状态失败: {str(db_e)}")
        finally:
            db.close()
            
        raise self.retry(exc=e)


@app.task(bind=True, queue='etf_data')
def schedule_daily_fund_update(self, fund_type="etf"):
    """
    调度每日基金净值数据更新任务
    每天收盘后执行，获取所有活跃基金的最新净值数据
    
    Args:
        fund_type: 基金类型，可选值：etf/lof/stock/bond/mix/money/qdii/index/all
    """
    logger.info(f"开始调度每日{fund_type}基金净值数据更新任务")
    
    try:
        # 首先更新基金列表
        update_fund_list.delay(fund_type)
        
        # 获取当前日期和前一交易日
        today = datetime.datetime.now().strftime('%Y%m%d')
        yesterday = (datetime.datetime.now() - datetime.timedelta(days=7)).strftime('%Y%m%d')
        # 注意：基金数据通常有延迟，所以获取7天前到今天的数据
        
        # 获取活跃基金
        db = next(get_db())
        query = db.query(SecurityInfo).filter(
            SecurityInfo.is_active == 1
        )
        
        # 根据基金类型筛选
        if fund_type == "all":
            query = query.filter(SecurityInfo.type.in_(["fund", "etf", "lof"]))
        elif fund_type == "etf":
            query = query.filter(SecurityInfo.type == "etf")
        elif fund_type == "lof":
            query = query.filter(SecurityInfo.type == "lof")
        else:
            query = query.filter(SecurityInfo.type == "fund")
        
        funds = query.all()
        db.close()
        
        logger.info(f"准备更新 {len(funds)} 只{fund_type}基金的净值数据")
        
        # 创建并行任务组
        tasks = []
        for fund in funds:
            tasks.append(collect_fund_nav_data.s(fund.code, yesterday, today))
        
        # 分批执行任务，避免一次创建过多任务
        batch_size = 50
        for i in range(0, len(tasks), batch_size):
            batch_tasks = tasks[i:i+batch_size]
            group(batch_tasks).apply_async()
            logger.info(f"已调度 {i+len(batch_tasks)}/{len(tasks)} 个基金净值数据采集任务")
        
        return {
            "status": "success",
            "message": f"已调度 {len(funds)} 只{fund_type}基金的净值数据更新任务"
        }
    except Exception as e:
        logger.error(f"调度{fund_type}基金更新任务失败: {str(e)}")
        raise


@app.task(bind=True, queue='etf_data')
def collect_fund_historical_data(self, fund_code, years=3):
    """
    采集基金的历史净值数据
    
    Args:
        fund_code: 基金代码
        years: 采集最近几年的数据，默认3年
    """
    logger.info(f"开始采集基金 {fund_code} 的 {years} 年历史净值数据")
    
    try:
        # 计算起始日期
        end_date = datetime.datetime.now().strftime('%Y%m%d')
        start_date = (datetime.datetime.now() - datetime.timedelta(days=365*years)).strftime('%Y%m%d')
        
        # 执行数据采集
        result = collect_fund_nav_data.delay(fund_code, start_date, end_date)
        
        return {
            "status": "success",
            "message": f"已调度基金 {fund_code} 的历史净值数据采集任务",
            "task_id": result.id
        }
    except Exception as e:
        logger.error(f"调度基金历史净值数据采集任务失败: {str(e)}")
        raise


# 测试函数
def test_fund_tasks():
    """测试基金任务函数"""
    # 更新基金列表
    update_result = update_fund_list.delay("etf")
    print(f"更新ETF基金列表任务已提交，任务ID: {update_result.id}")
    
    # 等待任务完成
    result = update_result.get(timeout=120)
    print(f"更新ETF基金列表结果: {result}")
    
    # 采集单只基金数据
    db = next(get_db())
    fund = db.query(SecurityInfo).filter(
        SecurityInfo.type.in_(["fund", "etf", "lof"])
    ).first()
    db.close()
    
    if fund:
        # 计算最近30天的日期范围
        end_date = datetime.datetime.now().strftime('%Y%m%d')
        start_date = (datetime.datetime.now() - datetime.timedelta(days=30)).strftime('%Y%m%d')
        
        collect_result = collect_fund_nav_data.delay(fund.code, start_date, end_date)
        print(f"基金 {fund.code} 净值数据采集任务已提交，任务ID: {collect_result.id}")
        
        # 等待任务完成
        result = collect_result.get(timeout=120)
        print(f"基金净值数据采集结果: {result}")
    else:
        print("没有找到基金数据，无法测试采集任务")


if __name__ == "__main__":
    # 配置日志
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
    )
    
    # 执行测试
    test_fund_tasks()